#include <iostream>
#include <eigen3/Eigen/Core>
#include <vector>
#include <functional>
#include <cmath>
typedef Eigen::Matrix<double, Eigen::Dynamic, 1> VectorXd;
typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> MatrixXd;
typedef std::function<double(double *)> Function;

enum Smoother
{
    WeightedJacobi = 1,
    GaussSeidel = 2
};

enum Restricter
{
    FullWeighting = 1,
    Injection = 2
};

enum Intper
{
    Linear = 1,
    Quadratic = 2
};

/**
 * @brief 多重网格求解器。
 */
class Multigrid_2D_Solver
{
protected:
    int _num_of_grid = 128;
    VectorXd _u_origin;
    VectorXd _f_origin;
    MatrixXd _BC;
    Smoother SmoothType = GaussSeidel;
    Restricter RestricType = FullWeighting;
    Intper IntpType = Linear;

public:
    VectorXd get_f_origin() { return _f_origin; };
    VectorXd get_u_origin() { return _u_origin; };
    MatrixXd get_BC() { return _BC; };
    double get_h() { return 1.0 / _num_of_grid; };
    void setSmoothType(int SmoothTypeCode) { SmoothType = (enum Smoother)SmoothTypeCode; };
    void setRestircType(int RestricTypeCode) { RestricType = (enum Restricter)RestricTypeCode; };
    void setIntpType(int IntpTypeCode) { IntpType = (enum Intper)IntpTypeCode; };
    void setNumOfGrid(int num) { _num_of_grid = num; };
    void Initialize_u();
    void InitializeRhs(Function _F);
    void Initialize(int num, Function _F, Function _BCF);
    VectorXd SmoothingforVC(VectorXd _u, VectorXd _f, double _h);
    VectorXd ResidualSolver(const VectorXd _u, const VectorXd _f, const double _h);
    VectorXd RestrictSolver(VectorXd _r);
    VectorXd IntpSolver(VectorXd _r);
    VectorXd Vcycle(VectorXd _u, VectorXd _f, double _h);
    VectorXd FMGcycle(VectorXd _f, double _h);
};

void Multigrid_2D_Solver::Initialize(int num, Function _F, Function _BCF)
{
    setNumOfGrid(num);
    _BC.setZero(num + 1, 4);
    double a[2] = {0, 0};
    ///第0列：边界y=0 第1列：边界y=1
    for (int j = 0; j <= 1; j++)
    {
        a[1] = j;
        for (int i = 0; i <= num; i++)
        {
            a[0] = 1.0 * i / num;
            _BC(i, j) = _BCF(a);
        };
    };
    ///第2列：边界x=0 第3列：边界x=1
    for (int j = 0; j <= 1; j++)
    {
        a[0] = j;
        for (int i = 0; i <= num; i++)
        {
            a[1] = 1.0 * i / num;
            _BC(i, j + 2) = _BCF(a);
        };
    };

    InitializeRhs(_F);
    Initialize_u();
    std::cout << "Initialized!" << std::endl;
};

void Multigrid_2D_Solver::Initialize_u()
{
    int N = _num_of_grid;
    _u_origin.setZero((1 + N) * (1 + N));
    for (int i = 0; i <= N; i++)
    {
        _u_origin(i) = _BC(i, 0);
        _u_origin((N + 1) * N + i) = _BC(i, 1);
        _u_origin(i * (N + 1)) = _BC(i, 2);
        _u_origin((i + 1) * (N + 1) - 1) = _BC(i, 3);
    };
};

void Multigrid_2D_Solver::InitializeRhs(Function _F)
{
    int N = _num_of_grid;
    double a[2] = {0, 0};
    _f_origin.setZero((N + 1) * (N + 1));
    for (int j = 1; j < N; j++)
    {
        a[1] = 1.0 * j / N;
        for (int i = 1; i < N; i++)
        {
            a[0] = 1.0 * i / N;
            _f_origin(j * (N + 1) + i) = _F(a);
            // std::cout<<j*(N+1)+i<<": "<< a[0]<< a[1]<<"  "<<_F(a) <<std::endl;
        };
    };
    for (int i = 0; i <= N; i++)
    {
        _f_origin(i) = _BC(i, 0);
        _f_origin((N + 1) * N + i) = _BC(i, 1);
        _f_origin(i * (N + 1)) = _BC(i, 2);
        _f_origin((i + 1) * (N + 1) - 1) = _BC(i, 3);
    };
};

VectorXd Multigrid_2D_Solver::SmoothingforVC(VectorXd _u, VectorXd _f, double _h)
{
    VectorXd _u_temp;
    int N = sqrt(_u.rows()) - 1;
    // std::cout << "We have N with " << N << std::endl;
    _u_temp.setZero((N + 1) * (N + 1));
    /// _u_temp 将继承 _u的所有ghost cell。
    for (int i = 0; i <= N; i++)
    {
        _u_temp(i) = _u(i);
        _u_temp((N + 1) * N + i) = _u((N + 1) * N + i);
        _u_temp(i * (N + 1)) = _u(i * (N + 1));
        _u_temp((i + 1) * (N + 1) - 1) = _u((i + 1) * (N + 1) - 1);
    };

    double _t = 0;
    double omega = 2.0 / 3;
    switch (SmoothType)
    {
    case GaussSeidel:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                _u_temp(j * (N + 1) + i) = (-_h * _h * _f(j * (N + 1) + i) + _u_temp((j - 1) * (N + 1) + i) + _u((j + 1) * (N + 1) + i) + _u_temp(j * (N + 1) + i - 1) + _u(j * (N + 1) + i + 1)) / 4.0;
            };
        };
        break;
    case WeightedJacobi:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                _t = (-_h * _h * _f(j * (N + 1) + i) + _u((j - 1) * (N + 1) + i) + _u((j + 1) * (N + 1) + i) + _u(j * (N + 1) + i - 1) + _u(j * (N + 1) + i + 1)) / 4.0;
                _u_temp(j * (N + 1) + i) = (1.0 - omega) * _u(j * (N + 1) + i) + omega * _t;
            };
        };
        break;
    default:
        break;
    };
    return _u_temp;
};

VectorXd Multigrid_2D_Solver::ResidualSolver(const VectorXd _u, const VectorXd _f, const double _h)
{
    VectorXd _u_temp;
    int N = sqrt(_u.rows()) - 1;
    _u_temp.setZero((N + 1) * (N + 1));
    for (int j = 1; j <= N - 1; j++)
    {
        for (int i = 1; i <= N - 1; i++)
        {
            _u_temp(j * (N + 1) + i) = _f(j * (N + 1) + i) - (_u((j - 1) * (N + 1) + i) + _u((j + 1) * (N + 1) + i) + _u(j * (N + 1) + i - 1) + _u(j * (N + 1) + i + 1) - 4.0 * _u(j * (N + 1) + i)) / (_h * _h);
        };
    };

    return _u_temp;
};

VectorXd Multigrid_2D_Solver::RestrictSolver(VectorXd _r)
{
    VectorXd _u_temp;
    int N = (sqrt(_r.rows()) - 1) / 2;
    _u_temp.setZero((N + 1) * (N + 1));
    double _temp = 0;
    switch (RestricType)
    {
    case FullWeighting:
        for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                _temp = (1.0 / 16.0) * (_r((2 * j - 1) * (2 * N + 1) + 2 * i - 1) + _r((2 * j - 1) * (2 * N + 1) + 2 * i + 1) + _r((2 * j + 1) * (2 * N + 1) + 2 * i - 1) + _r((2 * j + 1) * (2 * N + 1) + 2 * i + 1));
                _temp += (2.0 / 16.0) * (_r((2 * j - 1) * (2 * N + 1) + 2 * i) + _r((2 * j + 1) * (2 * N + 1) + 2 * i) + _r((2 * j) * (2 * N + 1) + 2 * i - 1) + _r((2 * j) * (2 * N + 1) + 2 * i + 1));
                _temp += (4.0 / 16.0) * (_r((2 * j) * (2 * N + 1) + 2 * i));
                _u_temp(j * (N + 1) + i) = _temp;
            };
        };
        break;
    case Injection:
        for (int i = 1; i <= N - 1; i++)
        {
            _u_temp(i) = _r(2 * i);
        };
        break;
    default:
        break;
    };

    ///ghost cell
    for (int i = 0; i <= N; i++)
    {
        _u_temp(i) = _r(2 * i);
        _u_temp((N + 1) * N + i) = _r((2 * N + 1) * 2 * N + 2 * i);
        _u_temp(i * (N + 1)) = _r(2 * i * (2 * N + 1));
        _u_temp((i + 1) * (N + 1) - 1) = _r((2 * i + 1) * (2 * N + 1) - 1);
    };

    return _u_temp;
};

VectorXd Multigrid_2D_Solver::IntpSolver(VectorXd _r)
{
    VectorXd _u_temp;
    int N = (sqrt(_r.rows()) - 1) * 2;
    _u_temp.setZero((N + 1) * (N + 1));
    if(_r(0)==_BC(0,0) && _r(_r.rows()-1)==_BC(_num_of_grid,1))
            {
                for (int i = 0; i <= N; i++)
            {
                _u_temp(i) = _BC(i, 0);
                _u_temp((N + 1) * N + i) = _BC(i, 1);
                _u_temp(i * (N + 1)) = _BC(i, 2);
                _u_temp((i + 1) * (N + 1) - 1) = _BC(i, 3);
            };
            };
    
    switch (IntpType)
    {
    case Linear:
         for (int j = 1; j <= N - 1; j++)
        {
            for (int i = 1; i <= N - 1; i++)
            {
                if (i%2==1 && j%2==1)
                {
                    _u_temp(j * (N + 1) + i) = 1.0/4*(_r((j-1)/2*(N/2+1)+(i-1)/2)+_r((j-1)/2*(N/2+1)+(i+1)/2)+_r((j+1)/2*(N/2+1)+(i-1)/2)+_r((j+1)/2*(N/2+1)+(i+1)/2));
                }
                else if(i%2==0 && j%2==0)
                {
                    _u_temp(j * (N + 1) + i) = _r(j/2*(N/2+1)+i/2);

                }
                else if(i%2==0 && j%2==1)
                {
                    _u_temp(j * (N + 1) + i) = 1.0/2*(_r((j+1)/2*(N/2+1)+i/2)+_r((j-1)/2*(N/2+1)+i/2));        
                }
                else if(i%2==1 && j%2==0)
                {
                    _u_temp(j * (N + 1) + i) = 1.0/2*(_r((j)/2*(N/2+1)+(i-1)/2)+_r((j)/2*(N/2+1)+(i+1)/2));
                };
            };
        };
        break;
    case Quadratic:
        break;
    default:
        break;
    };
    return _u_temp;
};

VectorXd Multigrid_2D_Solver::Vcycle(VectorXd _u, VectorXd _f, double _h)
{

    VectorXd __u, __r, __eps;
    __u = SmoothingforVC(_u, _f, _h);
    __r = ResidualSolver(__u, _f, _h);
    __r = RestrictSolver(__r);
    __eps.setZero(__r.rows());
    if ((sqrt(__eps.rows()) - 1) == 2)
    {
        __eps = SmoothingforVC(__eps, __r, 2 * _h);
    }
    else
    {
        __eps = Vcycle(__eps, __r, 2 * _h);
    };
    __u = __u + IntpSolver(__eps);
    __u = SmoothingforVC(__u, _f, _h);
    return __u;
};

VectorXd Multigrid_2D_Solver::FMGcycle(VectorXd _f, double _h)
{

    VectorXd __u, __r;
    int N = sqrt(_f.rows()) - 1;
    __u.setZero((N + 1)*(N+1));
    for (int i = 0; i <= N; i++)
    {
        __u(i) = _BC(i, 0);
        __u((N + 1) * N + i) = _BC(i, 1);
        __u(i * (N + 1)) = _BC(i, 2);
        __u((i + 1) * (N + 1) - 1) = _BC(i, 3);
    };

    if (N == 2)
    {
        __u = SmoothingforVC(__u, _f, _h);
    }
    else
    {
        __r = RestrictSolver(_f);
        __u = FMGcycle(__r, 2 * _h);
        __u = IntpSolver(__u);
        __u = Vcycle(__u, _f, _h);
    };
    return __u;
};